import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
from tensorflow import keras
import os

# acquire training and testing data
mnist = keras.datasets.mnist
(x_train, y_train), (x_test, y_test) = mnist.load_data()
num_classes, input_shape = 10, (28, 28, 1)

# format dataset
x_train, x_test = x_train / 255.0, x_test / 255.0
x_train = np.expand_dims(x_train, -1)
x_test = np.expand_dims(x_test, -1)
y_train = keras.utils.to_categorical(y_train, num_classes)
y_test = keras.utils.to_categorical(y_test, num_classes)

# set checkpoint
checkpoint_path = "Lab1 MNIST/training/checkpoint.ckpt"
checkpoint_dir = os.path.dirname(checkpoint_path)
cp_callback = keras.callbacks.ModelCheckpoint(filepath=checkpoint_path, save_weights_only=True, verbose=1)


def show_pic_examples():
    plt.figure(figsize=(10, 10))
    for i in range(25):
        plt.subplot(5, 5, i + 1)
        plt.xticks([])
        plt.yticks([])
        plt.imshow(x_train[i].squeeze())
        plt.xlabel(y_train[i].squeeze().argmax())
    plt.show()


def init_model():
    model = keras.Sequential([
        keras.Input(shape=input_shape),
        keras.layers.Conv2D(filters=32, kernel_size=5, activation='relu', padding='same'),
        keras.layers.MaxPooling2D(pool_size=2, padding='same'),
        keras.layers.Conv2D(filters=64, kernel_size=5, activation='relu', padding='same'),
        keras.layers.MaxPooling2D(pool_size=2, padding='same'),
        keras.layers.Flatten(),
        keras.layers.Dense(10, activation='softmax'),
    ])

    model.compile(optimizer='SGD', loss='categorical_crossentropy', metrics=['accuracy'])
    return model


def train_model():
    model = init_model()
    latest = tf.train.latest_checkpoint(checkpoint_dir)
    if latest is not None:
        model.load_weights(latest)
    model.fit(x_train, y_train, epochs=10, callbacks=[cp_callback])


def test_model():
    model = init_model()
    latest = tf.train.latest_checkpoint(checkpoint_dir)
    print(checkpoint_dir)
    if latest is not None:
        model.load_weights(latest)
    test_loss, test_acc = model.evaluate(x_test, y_test, verbose=1)
    print(test_acc)


def save_model():
    model = init_model()
    latest = tf.train.latest_checkpoint(checkpoint_dir)
    if latest is not None:
        model.load_weights(latest)
        model.save('Lab1 MNIST/model/trained_model')
        print('SUCCESS: MODEL SAVED TO [Lab1 MNIST/model/trained_model]')
    else:
        print('FAILURE: NO PARAMETERS FOUND')


if __name__ == '__main__':
    # train_model()
    # test_model()
    # save_model()
    show_pic_examples()
